{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import sys\n",
    "\n",
    "from collections import defaultdict\n",
    "from time import time\n",
    "from os import makedirs\n",
    "from os.path import join, isdir\n",
    "from glob import glob\n",
    "from keras.callbacks import TensorBoard\n",
    "\n",
    "sys.path.append('scripts')\n",
    "from model_helpers import *\n",
    "from data_helpers import *\n",
    "from train_helpers import *\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "os.environ['CUDA_VISIBLE_DEVICES']='1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IMGS_PATH = '/data'\n",
    "OUT_PATH = 'saved/logs'\n",
    "if not isdir(OUT_PATH):\n",
    "    makedirs(OUT_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "patch_size = (128, 128)\n",
    "center_size = (32, 32)\n",
    "num_train_test = num_test_test = 10\n",
    "batch_size = 128\n",
    "num_epochs = 500\n",
    "model_width = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "files = np.random.choice(glob(join(IMGS_PATH, '*jpg')), 10000, replace=False)\n",
    "train_files = files[:-(num_train_test + num_test_test)]\n",
    "train_test_files, test_test_files = files[-(num_train_test + num_test_test):-num_test_test], files[-num_test_test:]\n",
    "fake_files = glob('data/fake_files/*jpg')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imgss = defaultdict(list)\n",
    "start = time()\n",
    "for dataset, files in zip(['train', 'train_test', 'test_test', 'fake'], \n",
    "                          [train_files, train_test_files, test_test_files, fake_files]):\n",
    "    for i, f in enumerate(files):\n",
    "        imgss[dataset].append(cv2.imread(f, 0))\n",
    "        if i % 1000 == 0 and i != 0:\n",
    "            end = time()\n",
    "            print('Processing dataset {}: {:.0f} sec - ETA: {:.0f} sec'.format(\n",
    "                dataset, end-start, ((end-start) / (i + 1)) * (len(files) - i)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, len(imgss.keys()), figsize=(20, 30))\n",
    "for i, (dataset, imgs) in enumerate(imgss.items()):\n",
    "    ax[i].imshow(imgs[0], 'gray')\n",
    "    ax[i].set_title('{} image'.format(dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(30, 20))\n",
    "plt.imshow(imgss['fake'][1], 'gray')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Square images and resize to same size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for dataset, imgs in imgss.items():\n",
    "    res_imgs = [img[:min(img.shape), :min(img.shape)] for img in imgs] \n",
    "    min_size = np.min([img.shape[0] for img in res_imgs])\n",
    "\n",
    "    res_imgs = [cv2.resize(img, (min_size, min_size), interpolation=cv2.INTER_CUBIC) for img in res_imgs]\n",
    "    imgss[dataset] = np.expand_dims(res_imgs, axis=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create training generator and evaluation images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_generator = DataGenerator(imgss['train'], patch_size, center_size, batch_size=batch_size, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_test = []\n",
    "for patches in patchess.values():\n",
    "    x_test.extend(patches)\n",
    "x_test = np.array(x_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask = create_center_mask(patch_size, center_size)\n",
    "patchess = defaultdict(list)\n",
    "for i, (dataset, imgs) in enumerate(imgss.items()):\n",
    "    if dataset != 'train':\n",
    "        for j, img in enumerate(imgs):\n",
    "            patchess[dataset].append(prepare_patch(img, patch_size, mask, i + j))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(np.squeeze(patchess['fake'][3]), 'gray')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = create_anomaly_cnn(model_width=model_width)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.patch_size = patch_size\n",
    "model.center_size = center_size\n",
    "model.batch_size = batch_size\n",
    "model.num_epochs = num_epochs\n",
    "model.model_width = model_width"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.compile(optimizer='adam', loss=reconstruction_loss(patch_size, mask=mask), metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images_callback = TensorBoardImages(OUT_PATH, patchess, vis_every=1)\n",
    "checkpoint_callback = CustomModelCheckpoint(OUT_PATH, save_weights_only=False)\n",
    "losses_callback = TensorBoard(log_dir=OUT_PATH, batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "history_dict = model.fit_generator(train_generator, \n",
    "                                   validation_data=(x_test, x_test),\n",
    "                                   epochs=2, \n",
    "                                   verbose=1, \n",
    "                                   callbacks=[images_callback, checkpoint_callback, losses_callback], \n",
    "                                   workers=10, \n",
    "                                   use_multiprocessing=True, \n",
    "                                   shuffle=False, \n",
    "                                   initial_epoch=0)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (my-env)",
   "language": "python",
   "name": "myenv"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
